import cv2
import threading, time
import numpy as np
from xgolib import XGO
from PIL import Image,ImageDraw,ImageFont
import xgoscreen.LCD_2inch as LCD_2inch
import os
from ultralytics import YOLO
import math

class OptimizedPositionController():
    def __init__(self):
        # 初始化机械狗
        self.dog = XGO("xgolite")
        
        # 检查机械狗类型
        version = self.dog.read_firmware()
        if version[0] == 'M':
            print('检测到 XGO-MINI')
            self.dog = XGO("xgomini")
            self.dog_type = 'M'
        else:
            print('检测到 XGO-LITE')
            self.dog_type = 'L'
        
        # 初始化显示屏
        self.display = LCD_2inch.LCD_2inch()
        self.display.Init() 
        self.display.clear()
        self.splash = Image.new("RGB",(320,240),"black")
        self.display.ShowImage(self.splash)
        
        # 摄像头相关
        self.cap = None
        self.camera_still = False
        
        # 加载YOLO模型
        self.model = YOLO("best.pt")
        
        # 位置调整参数
        self.target_x = 160  # 理想夹取位置X坐标
        self.target_y = 180  # 理想夹取位置Y坐标
        self.tolerance = 15  # 位置容差（像素）
        
        # 移动参数（根据指导书优化）
        self.move_step = 8   # 移动步长
        self.turn_step = 20  # 旋转步长
        
        # 机械臂参数
        self.arm_x = 100  # 机械臂X位置
        self.arm_z = -30  # 机械臂Z位置
        
        # 状态变量
        self.is_adjusting = False
        self.last_detection = None
        self.adjustment_count = 0
        self.max_adjustments = 25
        
        # 初始化机械狗状态
        self.initialize_dog()
        
    def initialize_dog(self):
        """初始化机械狗状态"""
        print("🤖 初始化机械狗...")
        
        # 停止所有运动
        self.dog.stop()
        time.sleep(0.1)
        
        # 设置步频为正常
        self.dog.pace("normal")
        time.sleep(0.1)
        
        # 设置身高为正常
        self.dog.translation('z', 105)
        time.sleep(0.1)
        
        # 关闭自稳模式（便于手动控制）
        self.dog.imu(0)
        time.sleep(0.1)
        
        print("✅ 机械狗初始化完成")
    
    def open_camera(self):
        """打开摄像头"""
        if self.cap is None:
            self.cap = cv2.VideoCapture(0)
            self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320)
            self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 240)
            self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            self.cap.set(cv2.CAP_PROP_FPS, 30)
            print("✅ 摄像头已打开")
    
    def detect_target(self, frame):
        """检测长方体目标"""
        results = self.model.predict(frame, conf=0.3, verbose=False, imgsz=320)
        
        best_target = None
        best_conf = 0
        
        for r in results:
            if r.boxes is not None and len(r.boxes) > 0:
                for box in r.boxes:
                    conf = float(box.conf[0])
                    cls_id = int(box.cls[0])
                    label = self.model.names[cls_id]
                    
                    # 只检测长方体
                    if label == "长方体" and conf > best_conf:
                        best_conf = conf
                        x1, y1, x2, y2 = map(int, box.xyxy[0])
                        center_x = (x1 + x2) // 2
                        center_y = (y1 + y2) // 2
                        best_target = {
                            'center_x': center_x,
                            'center_y': center_y,
                            'box': (x1, y1, x2, y2),
                            'confidence': conf,
                            'width': x2 - x1,
                            'height': y2 - y1
                        }
        
        return best_target
    
    def calculate_adjustment(self, target):
        """计算需要的位置调整"""
        dx = target['center_x'] - self.target_x
        dy = target['center_y'] - self.target_y
        distance = math.sqrt(dx*dx + dy*dy)
        
        needs_adjustment = distance > self.tolerance
        
        # 确定调整方向
        if abs(dx) > 10:
            if dx > 0:
                x_direction = "right"  # 目标在右边，需要右移
            else:
                x_direction = "left"   # 目标在左边，需要左移
        else:
            x_direction = "center"
            
        if abs(dy) > 10:
            if dy > 0:
                y_direction = "backward"  # 目标在后面，需要后退
            else:
                y_direction = "forward"   # 目标在前面，需要前进
        else:
            y_direction = "center"
        
        return {
            'needs_adjustment': needs_adjustment,
            'distance': distance,
            'dx': dx,
            'dy': dy,
            'x_direction': x_direction,
            'y_direction': y_direction
        }
    
    def adjust_position(self, adjustment):
        """根据调整信息移动机械狗"""
        if not adjustment['needs_adjustment']:
            print("🎯 位置已对准！")
            return True
        
        print(f"🔧 调整位置: X误差={adjustment['dx']:.0f}, Y误差={adjustment['dy']:.0f}")
        
        # 优先调整X方向（左右移动）
        if adjustment['x_direction'] == "right":
            print("➡️ 向右移动")
            self.dog.move_y(-self.move_step)  # 使用move_y进行右移
            time.sleep(0.3)
            self.dog.move_y(0)  # 停止移动
        elif adjustment['x_direction'] == "left":
            print("⬅️ 向左移动")
            self.dog.move_y(self.move_step)   # 使用move_y进行左移
            time.sleep(0.3)
            self.dog.move_y(0)  # 停止移动
        
        # 然后调整Y方向（前后移动）
        elif adjustment['y_direction'] == "forward":
            print("⬆️ 向前移动")
            self.dog.move_x(self.move_step)   # 使用move_x进行前进
            time.sleep(0.3)
            self.dog.move_x(0)  # 停止移动
        elif adjustment['y_direction'] == "backward":
            print("⬇️ 向后移动")
            self.dog.move_x(-self.move_step)  # 使用move_x进行后退
            time.sleep(0.3)
            self.dog.move_x(0)  # 停止移动
        
        return False
    
    def prepare_grasp_pose(self):
        """准备夹取姿态"""
        print("🤖 调整夹取姿态...")
        
        # 蹲下 - 使用translation调整身高
        self.dog.translation('z', 85)
        time.sleep(0.5)
        
        # 调整头部角度 - 使用motor控制舵机
        self.dog.motor(11, -30)  # 头部向下
        self.dog.motor(21, -30)  # 头部向下
        time.sleep(0.5)
        
        # 准备机械臂 - 使用arm设置位置
        self.dog.arm(self.arm_x, self.arm_z)
        time.sleep(0.5)
        
        # 张开夹爪 - 使用claw控制
        self.dog.claw(0)
        time.sleep(0.5)
        
        print("✅ 夹取姿态准备完成")
    
    def execute_grasp(self):
        """执行夹取动作"""
        print("🦾 开始夹取...")
        
        # 闭合夹爪
        self.dog.claw(255)
        time.sleep(1)
        
        # 抬起机械臂
        self.dog.arm(self.arm_x, -10)
        time.sleep(1)
        
        # 恢复站立
        self.dog.translation('z', 105)
        time.sleep(1)
        
        # 复位所有状态
        self.dog.reset()
        time.sleep(2)
        
        print("✅ 夹取完成！")
    
    def auto_adjust_and_grasp(self):
        """自动调整位置并夹取"""
        print("🎯 启动自动位置调整和夹取模式")
        
        # 准备夹取姿态
        self.prepare_grasp_pose()
        
        # 开始检测和调整
        self.adjustment_count = 0
        
        while self.camera_still and self.adjustment_count < self.max_adjustments:
            # 获取图像
            ret, frame = self.cap.read()
            if not ret:
                print("无法读取摄像头图像")
                time.sleep(0.1)
                continue
            
            # 检测目标
            target = self.detect_target(frame)
            
            if target is not None:
                self.last_detection = target
                
                # 计算调整
                adjustment = self.calculate_adjustment(target)
                
                # 在图像上显示信息
                x1, y1, x2, y2 = target['box']
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                
                # 显示目标中心
                cv2.circle(frame, (target['center_x'], target['center_y']), 5, (0, 0, 255), -1)
                
                # 显示理想位置
                cv2.circle(frame, (self.target_x, self.target_y), 8, (255, 0, 0), 2)
                
                # 显示连接线
                cv2.line(frame, (target['center_x'], target['center_y']), 
                        (self.target_x, self.target_y), (255, 255, 0), 2)
                
                # 显示调整信息
                info_text = f"DX:{adjustment['dx']:.0f} DY:{adjustment['dy']:.0f}"
                cv2.putText(frame, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
                
                # 显示调整次数
                count_text = f"Adjust: {self.adjustment_count}/{self.max_adjustments}"
                cv2.putText(frame, count_text, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
                
                # 执行位置调整
                if adjustment['needs_adjustment']:
                    if self.adjust_position(adjustment):
                        print("🎉 位置调整完成，开始夹取！")
                        time.sleep(1)  # 等待稳定
                        self.execute_grasp()
                        return True
                    self.adjustment_count += 1
                else:
                    print("🎯 位置已对准，开始夹取！")
                    time.sleep(1)  # 等待稳定
                    self.execute_grasp()
                    return True
            else:
                # 没有检测到目标
                cv2.putText(frame, "No target", (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
            
            # 显示图像
            try:
                image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                image_rgb = cv2.flip(image_rgb, 1)
                pil_image = Image.fromarray(image_rgb)
                self.display.ShowImage(pil_image)
            except Exception as e:
                print(f"显示图像错误: {e}")
            
            time.sleep(0.1)
        
        print("⚠️ 调整次数过多或未检测到目标")
        return False
    
    def camera_mode(self):
        """摄像头模式 - 实时显示检测结果"""
        if not self.cap.isOpened():
            print("无法打开摄像头")
            return
        
        print("📷 摄像头模式启动")
        
        while self.camera_still and self.cap is not None and self.cap.isOpened():
            ret, frame = self.cap.read()
            if not ret:
                print("无法读取摄像头图像")
                time.sleep(0.1)
                continue
            
            # 检测目标
            target = self.detect_target(frame)
            
            if target is not None:
                # 计算调整信息
                adjustment = self.calculate_adjustment(target)
                
                # 绘制检测框
                x1, y1, x2, y2 = target['box']
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                
                # 绘制中心点
                cv2.circle(frame, (target['center_x'], target['center_y']), 5, (0, 0, 255), -1)
                
                # 绘制目标位置
                cv2.circle(frame, (self.target_x, self.target_y), 8, (255, 0, 0), 2)
                
                # 显示信息
                info_text = f"DX:{adjustment['dx']:.0f} DY:{adjustment['dy']:.0f}"
                cv2.putText(frame, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
                
                if adjustment['needs_adjustment']:
                    status_text = "需要调整"
                    color = (0, 0, 255)
                else:
                    status_text = "位置正确"
                    color = (0, 255, 0)
                
                cv2.putText(frame, status_text, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
                
                # 显示置信度
                conf_text = f"Conf: {target['confidence']:.2f}"
                cv2.putText(frame, conf_text, (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
            else:
                cv2.putText(frame, "No target", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
            
            # 显示图像
            try:
                image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                image_rgb = cv2.flip(image_rgb, 1)
                pil_image = Image.fromarray(image_rgb)
                self.display.ShowImage(pil_image)
            except Exception as e:
                print(f"显示图像错误: {e}")
            
            time.sleep(0.05)
    
    def manual_control(self):
        """手动控制模式"""
        print("🔧 手动控制模式")
        print("命令: w(前进) s(后退) a(左移) d(右移) q(左转) e(右转)")
        print("      z(蹲下) x(站起) g(夹取) r(复位) c(退出)")
        print("      h(升高) l(降低) t(抬头) b(低头)")
        
        while True:
            try:
                key = input("请输入命令: ").lower().strip()
                
                if key == 'w':
                    self.dog.move_x(10)
                    time.sleep(0.3)
                    self.dog.move_x(0)
                elif key == 's':
                    self.dog.move_x(-10)
                    time.sleep(0.3)
                    self.dog.move_x(0)
                elif key == 'a':
                    self.dog.move_y(10)
                    time.sleep(0.3)
                    self.dog.move_y(0)
                elif key == 'd':
                    self.dog.move_y(-10)
                    time.sleep(0.3)
                    self.dog.move_y(0)
                elif key == 'q':
                    self.dog.turn(20)
                    time.sleep(0.3)
                    self.dog.turn(0)
                elif key == 'e':
                    self.dog.turn(-20)
                    time.sleep(0.3)
                    self.dog.turn(0)
                elif key == 'z':
                    self.dog.translation('z', 85)
                elif key == 'x':
                    self.dog.translation('z', 105)
                elif key == 'h':
                    self.dog.translation('z', 115)
                elif key == 'l':
                    self.dog.translation('z', 75)
                elif key == 't':
                    self.dog.motor(11, 30)
                    self.dog.motor(21, 30)
                elif key == 'b':
                    self.dog.motor(11, -30)
                    self.dog.motor(21, -30)
                elif key == 'g':
                    self.execute_grasp()
                elif key == 'r':
                    self.dog.reset()
                elif key == 'c':
                    break
                else:
                    print("无效命令")
                    
            except KeyboardInterrupt:
                break
        
        print("手动控制模式结束")
    
    def start_camera(self):
        """启动摄像头"""
        self.open_camera()
        if self.cap and self.cap.isOpened():
            self.camera_still = True
            t = threading.Thread(target=self.camera_mode, daemon=True)
            t.start()
            
            # 调整摄像头角度
            self.dog.motor(11, -30)
            self.dog.motor(21, -30)
            print("✅ 摄像头启动完成")
        else:
            print("❌ 摄像头启动失败")
    
    def stop_camera(self):
        """停止摄像头"""
        self.camera_still = False
        time.sleep(0.2)
        if self.cap:
            self.cap.release()
            self.cap = None
        splash = Image.new("RGB", (320, 240), "black")
        self.display.ShowImage(splash)
        print("✅ 摄像头已关闭")
    
    def close(self):
        """清理资源"""
        self.stop_camera()
        self.display.clear()
        print("💤 系统已关闭")

# 主程序
if __name__ == "__main__":
    print("🤖 机械狗优化位置调整系统")
    print("=" * 50)
    print("基于官方接口函数优化")
    print("功能:")
    print("1. 实时检测长方体位置")
    print("2. 自动调整机械狗位置")
    print("3. 自动夹取长方体")
    print("4. 手动控制调试")
    print("=" * 50)
    
    controller = OptimizedPositionController()
    
    try:
        print("请选择模式:")
        print("1. 实时检测模式")
        print("2. 自动调整夹取模式")
        print("3. 手动控制模式")
        
        choice = input("请输入选择 (1-3): ").strip()
        
        if choice == "1":
            # 实时检测模式
            controller.start_camera()
            print("📷 实时检测模式启动，按 Ctrl+C 退出")
            while True:
                time.sleep(0.1)
                
        elif choice == "2":
            # 自动调整夹取模式
            controller.start_camera()
            time.sleep(2)  # 等待摄像头稳定
            success = controller.auto_adjust_and_grasp()
            if success:
                print("🎉 自动夹取成功！")
            else:
                print("❌ 自动夹取失败")
                
        elif choice == "3":
            # 手动控制模式
            controller.manual_control()
            
        else:
            print("无效选择，启动实时检测模式")
            controller.start_camera()
            while True:
                time.sleep(0.1)
                
    except KeyboardInterrupt:
        print("\n⚠️ 用户中断程序")
    finally:
        controller.close()
        print("🎯 程序结束") 